#ifndef _BUOLA_CV_TREES_CFLANNTREE_H_
#define _BUOLA_CV_TREES_CFLANNTREE_H_

#include <buola/mat.h>
#include <buola/cv/ccloud.h>

#include <flann/flann.hpp>

namespace buola { namespace cvi {
    
template<typename tScalar>
class CFlannTree
{
    typedef flann::Index<flann::L2_Simple<tScalar>> TIndex;
    
public:
    CFlannTree()
    {}

    template<typename tPoint>
    CFlannTree(const CCloud<tPoint> &pCloud)
    {
        SetContent(pCloud);
    }

    //not copying the data into a temporary actually makes it slower!
    template<typename tPoint>
    void SetContent(const CCloud<tPoint> &pCloud)
    {
        mData.resize(pCloud.size()*3);
        tScalar *lDataPtr=mData.data();
        for(int i=0;i<pCloud.size();i++)
        {
            *lDataPtr++=pCloud[i].mPoint.x();
            *lDataPtr++=pCloud[i].mPoint.y();
            *lDataPtr++=pCloud[i].mPoint.z();
        }
        
        flann::Matrix<tScalar> lDataMat(mData.data(),pCloud.size(),3);
        mFlann.reset(new TIndex(lDataMat,flann::KDTreeSingleIndexParams(15)));
        start_timer();
        mFlann->buildIndex();
        end_timer();
    }
    
    template<typename tOScalar>
    std::vector<int> GetNearestK(const mat::CVec<tOScalar,3> &pPoint,int k) const
    {
        tScalar lPoint[3];
        lPoint[0]=pPoint.x();
        lPoint[1]=pPoint.y();
        lPoint[2]=pPoint.z();
        std::vector<int> lIndices(k);
        std::vector<tScalar> lDistances(k);
        flann::Matrix<tScalar> lPointMat(lPoint,1,3);
        flann::Matrix<int> lIndicesMat(lIndices.data(),1,k);
        flann::Matrix<tScalar> lDistancesMat(lDistances.data(),1,k);
        mFlann->knnSearch(lPointMat,lIndicesMat,lDistancesMat,k,flann::SearchParams(-1));
        return lIndices;
    }
    
    template<typename tOScalar>
    std::pair<int,tScalar> GetNearest(const mat::CVec<tOScalar,3> &pPoint) const
    {
        tScalar lPoint[3];
        lPoint[0]=pPoint.x();
        lPoint[1]=pPoint.y();
        lPoint[2]=pPoint.z();
        int lIndex;
        tScalar lDistance;
        flann::Matrix<tScalar> lPointMat(lPoint,1,3);
        flann::Matrix<int> lIndicesMat(&lIndex,1,1);
        flann::Matrix<tScalar> lDistancesMat(&lDistance,1,1);
        mFlann->knnSearch(lPointMat,lIndicesMat,lDistancesMat,1,flann::SearchParams(-1));
        return {lIndex,lDistance};
    }
    
    template<typename tOScalar>
    std::vector<std::pair<int,tScalar>> GetNearest(const std::vector<mat::CVec<tOScalar,3>> &pPoints) const
    {
        std::vector<tScalar> lData(pPoints.size()*3);
        tScalar *lDataPtr=lData.data();
        for(int i=0;i<pPoints.size();i++)
        {
            *lDataPtr++=pPoints[i].x();
            *lDataPtr++=pPoints[i].y();
            *lDataPtr++=pPoints[i].z();
        }
        std::vector<int> lIndices(pPoints.size());
        std::vector<tScalar> lDistances(pPoints.size());
        flann::Matrix<tScalar> lPointsMat(lData.data(),pPoints.size(),3);
        flann::Matrix<int> lIndicesMat(lIndices.data(),pPoints.size(),1);
        flann::Matrix<tScalar> lDistancesMat(lDistances.data(),pPoints.size(),1);

        mFlann->knnSearch(lPointsMat,lIndicesMat,lDistancesMat,1,flann::SearchParams(-1));

        std::vector<std::pair<int,tScalar>> lResult;
        
        for(int i=0;i<pPoints.size();i++)
        {
            lResult.emplace_back(lIndices[i],lDistances[i]);
        }

        return lResult;
    }
    
    mat::CVec<tScalar,3> operator[](int i) const
    {
        return {mData[i*3],mData[i*3+1],mData[i*3+2]};
    }
    
private:
    std::shared_ptr<TIndex> mFlann;
    std::vector<tScalar> mData;
};
    
template<typename tPoint>
std::vector<std::pair<int,int>> get_correspondences(const CCloud<tPoint> &pCloud,const CFlannTree<double> &pTree,
                                                    double pMaxDist)
{
    std::vector<std::pair<int,int>> lCorrespondences;
    lCorrespondences.reserve(pCloud.size());
    
    double lMaxDistSq=sq(pMaxDist);
    
    std::vector<mat::CVec3d> lPoints(pCloud.size());
    
    for(int i=0;i<pCloud.size();i++)
        lPoints[i]=pCloud[i].mPoint;

    std::vector<std::pair<int,double>> lMatches=pTree.GetNearest(lPoints);

    for(int i=0;i<lMatches.size();i++)
    {
        if(lMatches[i].second>lMaxDistSq) continue;
        lCorrespondences.emplace_back(i,lMatches[i].first);
    }
    
    return lCorrespondences;
}

/*namespace cvi*/ } /*namespace buola*/ }

#endif
